from csp import backtracking_search, CSP
import itertools

def restricciones(A, a, B, b):
    "A constraint saying two neighboring variables must differ in value."
    def checkop(Aux):
        if Aux['R'] == Aux['O'] + Aux['O'] - Aux['C1']*10:
            if Aux['U'] == Aux['W'] + Aux['W'] + Aux['C1'] - Aux['C2']*10:
                if Aux['O'] == Aux['T'] + Aux['T'] + Aux['C2'] - Aux['C3']*10:
                    if Aux['F'] == Aux['C3']:
                        return True
        return False


    if A =='AUX':
        if B =='AUX':
            if a == b:
                checkop(a)
        else:
            return a[B] == b and checkop(a)
    if B =='AUX':
        return b[A] == a and checkop(b)
    else:
        return a != b



def incluir(lista,numero):
    if (numero in lista): return False
    else:
        lista.append(numero)
        return True



def CriptoAritmeticaCSP(letras,numeros):

    variables=letras

    #paso numeros a diccionario
    dominio={}
    for i in letras:
        dominio.update({i:numeros})

    variables.append('C1')
    variables.append('C2')
    variables.append('C3')
    dominio.update({'C1':[0,1]})
    dominio.update({'C2':[0,1]})
    dominio.update({'C3':[0,1]})

    lista=[]
    for element in itertools.product(dominio['F'],dominio['O'],dominio['U'],dominio['R'],dominio['W'],
                                     dominio['T'],dominio['C1'],dominio['C2'],dominio['C3']):
        l=[]        
        aux={}
        bandera = 1
        
        aux.update({'F':element[0]})
        if not incluir(l,element[0]):
            bandera=0
        aux.update({'O':element[1]})
        if not incluir(l,element[1]):
            bandera=0
        aux.update({'U':element[2]})
        if not incluir(l,element[2]):
            bandera=0
        aux.update({'R':element[3]})
        if not incluir(l,element[3]):
            bandera=0
        aux.update({'T':element[4]})
        if not incluir(l,element[4]):
            bandera=0
        aux.update({'W':element[5]})
        if not incluir(l,element[5]):
            bandera=0


        aux.update({'C1':element[6]})
        aux.update({'C2':element[7]})
        aux.update({'C3':element[8]})

        

        if aux['C1'] == 1:
            if aux['O']*2 < 10:
                bandera = 0
        if aux['C2'] == 1:
            if aux['W']*2 + aux['C1'] < 10:
                bandera = 0
        if aux['C3'] == 1:
            if aux['T']*2 + aux['C2'] < 10:
                bandera = 0

        if bandera==1:
            lista.append(aux)

    variables.append('AUX')
    dominio.update({'AUX':lista})

    vecinos = {'T':['W','O','F','U','R','C1','C2','C3','AUX'],'W':['T','O','F','U','R','C1','C2','C3','AUX'],'O':['T','W','F','U','R','C1','C2','C3','AUX'],
                'F':['T','W','O','U','R','C1','C2','C3','AUX'],'U':['T','W','O','F','R','C1','C2','C3','AUX'],'R':['T','W','O','F','U','C1','C2','C3','AUX'],
                'C1':['T','W','O','F','U','R','C2','C3','AUX'],'C2':['T','W','O','F','U','R','C1','C3','AUX'],'C3':['T','W','O','F','U','R','C1','C2','AUX'],
                'AUX':['T','W','O','F','U','R','C1','C2','C3']}

    return CSP(variables, dominio, vecinos, restricciones)

le = ['T','W','O','F','U','R']
nu = [1,2,3,4,5,6,7,8,9,0]
#twoplustwo= CriptoAritmeticaCSP(le,nu)

print backtracking_search(CriptoAritmeticaCSP(le,nu),fc=True)

